#include "pch.h"
#include "RegSyncRoot.h"
#include "GaojieDriveApp.h"
#include "AppConfig.h"
#include "..\Common\StrHelper.h"

namespace
{

std::wstring GetWindowsSid()
{
	char szBuffer[256];
	DWORD dwLength = sizeof(szBuffer);
	std::shared_ptr<void> pTokenInfo(szBuffer, [](void *) {});

	HANDLE hToken = GetCurrentProcessToken();

	if (!GetTokenInformation(hToken, TokenUser, pTokenInfo.get(), dwLength, &dwLength))
	{
		if (GetLastError() != ERROR_INSUFFICIENT_BUFFER)
			return L"";

		pTokenInfo.reset(malloc(dwLength), free);

		if (!GetTokenInformation(hToken, TokenUser, pTokenInfo.get(), dwLength, &dwLength))
			return L"";
	}

	std::wstring strWindowsSid;
	LPWSTR lpStringSid = nullptr;

	if (ConvertSidToStringSid(std::static_pointer_cast<TOKEN_USER>(pTokenInfo)->User.Sid, &lpStringSid))
	{
		strWindowsSid = lpStringSid;
		LocalFree(lpStringSid);
	}

	return strWindowsSid;
}

std::wstring GetSyncRootId(const std::wstring& strWindowsSid)
{
	std::wstring strSyncRootId;
	Str::Format(strSyncRootId, L"GaojieDrive!%s!Personal", strWindowsSid);
	return strSyncRootId;
}

void SetRegString(HKEY hKey, LPCWSTR lpName, LPCWSTR lpValue, size_t nValueLen, int (*pfCompare)(LPCWSTR, LPCWSTR))
{
	if (nValueLen == -1)
		nValueLen = wcslen(lpValue);

	DWORD dwType, dwSize = (DWORD)(nValueLen + 1) * sizeof(wchar_t);
	LPWSTR lpBuffer = (LPWSTR)alloca(dwSize);

	if (RegQueryValueExW(hKey, lpName, NULL, &dwType, (BYTE *)lpBuffer, &dwSize) != ERROR_SUCCESS ||
		dwType != REG_SZ ||
		pfCompare(lpValue, lpBuffer) != 0)
	{
		RegSetValueExW(hKey, lpName, 0, REG_SZ, (const BYTE *)lpValue, dwSize);
	}
}

void SetDisplayName(HKEY hKey)
{
	SetRegString(hKey, L"DisplayNameResource", L"GaojieDrive", -1, wcscmp);
}

void SetIconPath(HKEY hKey)
{
	wchar_t szIconPath[MAX_PATH + 2] = {};
	if (GetModuleFileNameW(NULL, szIconPath, MAX_PATH))
	{
		wcscat_s(szIconPath, L",0");
		SetRegString(hKey, L"IconResource", szIconPath, -1, _wcsicmp);
	}
}

void SetSyncRootPath(HKEY hKey, const std::wstring& strValueName)
{
	HKEY hKey2;
	if (RegCreateKeyExW(hKey, L"UserSyncRoots", 0, NULL, REG_OPTION_NON_VOLATILE, KEY_QUERY_VALUE | KEY_SET_VALUE, NULL, &hKey2, NULL) == ERROR_SUCCESS)
	{
		std::wstring strRootPath = GetApp()->GetAppConfig()->GetClientRoot();
		if (strRootPath.back() == '\\')
			strRootPath.pop_back();

		SetRegString(hKey2, strValueName.c_str(), strRootPath.c_str(), strRootPath.size(), _wcsicmp);
		RegCloseKey(hKey2);
	}
}

void AddSyncRootManager()
{
	std::wstring strWindowsSid = GetWindowsSid();
	std::wstring strSubKey = LR"(Software\Microsoft\Windows\CurrentVersion\Explorer\SyncRootManager\)" + GetSyncRootId(strWindowsSid);

	HKEY hKey;
	if (RegCreateKeyExW(HKEY_LOCAL_MACHINE, strSubKey.c_str(), 0, NULL, REG_OPTION_NON_VOLATILE, KEY_CREATE_SUB_KEY | KEY_QUERY_VALUE | KEY_SET_VALUE, NULL, &hKey, NULL) == ERROR_SUCCESS)
	{
		SetDisplayName(hKey);
		SetIconPath(hKey);
		DWORD dwFlags = 0x62;
		RegSetValueExW(hKey, L"Flags", 0, REG_DWORD, (const BYTE *)&dwFlags, sizeof(dwFlags));
		SetSyncRootPath(hKey, strWindowsSid);
		RegCloseKey(hKey);
	}
}

}	// namespace

void RegisterSyncRoot()
{
	AddSyncRootManager();
}
